import os
import stat
import mmap
import struct
from binascii import hexlify
import hashlib
# Constant separating blocks in the .blk files
BITCOIN_CONSTANT = b"\xf9\xbe\xb4\xd9"

BlocksPath = r"E:\OneDrive - 西南大学\项目资料\202103 与南瑞的合作-区块链一网通办\17-论文-快速检索索引\03-Program\blocks"
def get_files(path):
    if not stat.S_ISDIR(os.stat(path)[stat.ST_MODE]):
        return [path]
    files = os.listdir(path)
    files = [f for f in files if f.startswith("blk") and f.endswith(".dat")]
    files = map(lambda x: os.path.join(path, x), files)
    return sorted(files)

def get_blocks(blockfile):
    """
    Given the name of a .blk file, for every block contained in the file,
    yields its raw hexadecimal value
    """
    with open(blockfile, "rb") as f:
        if os.name == 'nt':
            size = os.path.getsize(f.name)
            raw_data = mmap.mmap(f.fileno(), size, access=mmap.ACCESS_READ)
        else:
            # Unix-only call, will not work on Windows, see python doc.
            raw_data = mmap.mmap(f.fileno(), 0, prot=mmap.PROT_READ)
        length = len(raw_data)
        offset = 0
        block_count = 0
        while offset < (length - 4):
            if raw_data[offset:offset+4] == BITCOIN_CONSTANT:
                offset += 4
                size = struct.unpack("<I", raw_data[offset:offset+4])[0]
                offset += 4 + size
                block_count += 1
                yield raw_data[offset-size:offset]
            else:
                offset += 1
        raw_data.close()

def format_hash(hash_):
    return str(hexlify(hash_[::-1]).decode("utf-8"))


def double_sha256(data):
    return hashlib.sha256(hashlib.sha256(data).digest()).digest()


def decode_varint(data):
    assert(len(data) > 0)
    size = int(data[0])
    assert(size <= 255)

    if size < 253:
        return size, 1

    if size == 253:
        format_ = '<H'
    elif size == 254:
        format_ = '<I'
    elif size == 255:
        format_ = '<Q'
    else:
        # Should never be reached
        assert 0, "unknown format_ for size : %s" % size

    size = struct.calcsize(format_)
    return struct.unpack(format_, data[1:size+1])[0], size + 1

    
def get_block_transactions(raw_hex):
    """Given the raw hexadecimal representation of a block,
    yields the block's transactions
    """
    # Skipping the header
    transaction_data = raw_hex[80:]

    # Decoding the number of transactions, offset is the size of
    # the varint (1 to 9 bytes)
    n_transactions, offset = decode_varint(transaction_data)

    for i in range(n_transactions):
        # Try from 1024 (1KiB) -> 1073741824 (1GiB) slice widths
        for j in range(0, 20):
            try:
                offset_e = offset + (1024 * 2 ** j)
                transaction = Transaction.from_hex(
                    transaction_data[offset:offset_e])
                yield transaction
                break
            except:
                continue

        # Skipping to the next transaction
        offset += transaction.size



class Block(object):
    """
    Represents a Bitcoin block, contains its header and its transactions.
    """

    def __init__(self, raw_hex, height=None, blk_file=None):
        self.hex = raw_hex
        self._hash = None
        self._transactions = None
        self._header = None
        self._n_transactions = None
        self.size = len(raw_hex)
        self.height = height
        self.blk_file = blk_file

    def __repr__(self):
        return "Block(%s)" % self.hash

    @classmethod
    def from_hex(cls, raw_hex):
        """Builds a block object from its bytes representation"""
        return cls(raw_hex)

    @property
    def hash(self):
        """Returns the block's hash (double sha256 of its 80 bytes header"""
        if self._hash is None:
            self._hash = format_hash(double_sha256(self.hex[:80]))
        return self._hash

    @property
    def n_transactions(self):
        """Return the number of transactions contained in this block,
        it is faster to use this than to use len(block.transactions)
        as there's no need to parse all transactions to get this information
        """
        if self._n_transactions is None:
            self._n_transactions = decode_varint(self.hex[80:])[0]

        return self._n_transactions

    @property
    def transactions(self):
        """Returns a list of the block's transactions represented
        as Transaction objects"""
        if self._transactions is None:
            self._transactions = list(get_block_transactions(self.hex))

        return self._transactions

    @property
    def header(self):
        """Returns a BlockHeader object corresponding to this block"""
        if self._header is None:
            self._header = BlockHeader.from_hex(self.hex[:80])
        return self._header


for blk_file in get_files(BlocksPath):
    print(blk_file)
    for raw_block in get_blocks(blk_file):
        print(raw_block)
        yield Block(raw_block, None, os.path.split(blk_file)[1])